'''
Created on Jun 21, 2011

@author: Nam Khanh Tran, Ba Dat Nguyen
'''

import sys
import feature_selector

def predict(weight, x_doc):
    """
    """
    pre_cat = None
    max_sum = -1
    
    for cat in weight.keys():
        sum = 0
        for word in x_doc:
            sum = sum + weight[cat][word]
        if sum > max_sum:
            max_sum = sum
            pre_cat = cat
    
    return pre_cat

def update(pre_cat, cat, weight, x_doc):
    """
    """
    for word in x_doc:
        weight[cat][word] = weight[cat][word] + 1
        weight[pre_cat][word] = weight[pre_cat][word] - 1
    
    return weight

def special_modeling(iterN, word_list, categories, documents):
    """
    """
    weight = dict()
    for cat in categories.keys():
        weight[cat] = dict()
        for word in word_list:
            weight[cat][word] = 0
            
    for i in range(0, iterN):
        print "Processing iteration k = ", i+1
        for cat in categories.keys():
            doc_list = categories[cat]
            for doc in doc_list:
                pred_cat = predict(weight, documents[doc])
                if pred_cat != cat:
                    weight = update(pred_cat, cat, weight, documents[doc])             
                # end if
            # end for
        # end for
    print "Done."
    return weight
    

def modeling(iterN, word_list, categories, documents):
    """
    """
    weight = dict()
    for cat in categories.keys():
        weight[cat] = dict()
        for word in word_list:
            weight[cat][word] = 0
    
    print "Pre-Processing..."
    word_set = set(word_list)
    x = dict()
    for doc in documents.keys():
        temp = list()
        for word in documents[doc]:
            if word in word_set:
                temp.append(word)
        x[doc] = temp
        
    for i in range(0, iterN):
        print "Processing iteration k = ", i+1        
        for cat in categories.keys():
            doc_list = categories[cat]
            for doc in doc_list:
                pred_cat = predict(weight, x[doc])
                if pred_cat != cat:
                    weight = update(pred_cat, cat, weight, x[doc])
             
                # end if
            # end for
        # end for
    print "Done."
    return weight

def train_model(iterN = 10, topN = -1, method = -1):
    """
    """
    print 'Training...'
    dict_path = "20-ng/train"
    (vocab, categories, vocab_cat, documents) = feature_selector.read_file(dict_path)
    
    if topN == -1:
        return special_modeling(iterN, vocab.keys(), categories, documents)
    
    # information_gain
    if method == 1:
        #info_gain = feature_selector.info_gain(vocab, categories, vocab_cat)
        print "information gain"
        word_list = list()
        with open('ig_top' + str(topN) + '.txt', 'r') as fin:
            for word in fin.read().split('\n'):
                word_list.append(word)
            
        return modeling(iterN, word_list, categories, documents)            
    # chi_square
    else:
        print "chi_square"
        #chi_square = feature_selector.chi_square(vocab, categories, vocab_cat)
        word_list = list()
        with open('cs_top' + str(topN) + '.txt', 'r') as fin:
            for word in fin.read().split('\n'):
                word_list.append(word)
                           
        return modeling(iterN, word_list, categories, documents)   

def testing(weight, word_list, categories, documents):
    """
    """
    word_set = set(word_list)
    x = dict()
    for doc in documents.keys():
        temp = list()
        for word in documents[doc]:
            if word in word_set:
                temp.append(word)
        x[doc] = temp

    count = 0
    
    for cat in categories.keys():
        doc_list = categories[cat]
        for doc in doc_list:
            pred_cat = predict(weight, x[doc])
            if pred_cat == cat:
                count = count + 1
                
    print count
    print 'Accuracy = ', count * 100.0 / len(documents) 

def test_model(weight, topN = -1, method = 1):
    """
    """
    print "Testing...."
    dict_path = "20-ng/test"
    (vocab, categories, vocab_cat, documents) = feature_selector.read_file(dict_path)
    
    if topN == -1:
        word_list = list()
        with open('training_vocab.txt', 'r') as fin:
            for word in fin.read().split('\n'):
                word_list.append(word)
        
        return testing(weight, word_list, categories, documents)
    
    # information_gain
    if method == 1:
        #info_gain = feature_selector.info_gain(vocab, categories, vocab_cat)
        word_list = list()
        with open('ig_top' + str(topN) + '.txt', 'r') as fin:
            for word in fin.read().split('\n'):
                word_list.append(word)
            
        return testing(weight, word_list, categories, documents)            
    # chi_square
    else:
        #chi_square = feature_selector.chi_square(vocab, categories, vocab_cat)
        word_list = list()
        with open('cs_top' + str(topN) + '.txt', 'r') as fin:
            for word in fin.read().split('\n'):
                word_list.append(word)
                           
        return testing(weight, word_list, categories, documents)    
    
if __name__ == "__main__":
    if len(sys.argv) < 4:
        print '\nUsage: python {0} t [-1 100 1000] [1 2]'.format(sys.argv[0])
        print 't : number of iterations'
        print '[-1 100 1000]: topN features. -1 for using all features'
        print '[1 2]: infomation gain or chi square'
        sys.exit(1)
    
    weight = train_model(iterN = int(sys.argv[1]), topN = int(sys.argv[2]), method=int(sys.argv[3]))
    test_model(weight, topN = int(sys.argv[2]), method = int(sys.argv[3]))    